feat: implement Kafka-to-Spark ingestion pipeline (Chapter 4) - #3

🤖 AI PR Review Summary

This Pull Request introduces a new "Ingestion" project, demonstrating real-time event ingestion using Kafka and PySpark. It includes a Kafka producer, a PySpark job to consume and process JSON events, and local Docker Compose setup for Kafka, Zookeeper, and Schema Registry. The project also updates README.md, agents.md, and requirements.txt to reflect the new structure and dependencies.

The documentation and overall setup are excellent, providing clear instructions and a functional example for local development. However, there's a significant discrepancy between the project's stated goal of "real-time event ingestion using Kafka and Spark Streaming" and the actual implementation of the Spark job, which performs a batch read from Kafka rather than a continuous stream.

💡 Key Feedback & Recommendations

1. Major: Streaming vs. Batch Implementation Mismatch

Issue: The README.md and the docstring for kafka_json_to_file_job.py describe the job as "Spark Streaming ingestion job" and "real-time event ingestion." However, the PySpark job uses spark.read.format("kafka").options(...).load(), which performs a batch read of all available messages up to the point of execution (based on startingOffsets). For true streaming or real-time ingestion, spark.readStream and df.writeStream should be used.

Impact: The current implementation does not demonstrate continuous, real-time processing as suggested by the documentation. It's a one-time historical load.

Recommendation: To align with the "streaming" description, the PySpark job should be converted to a streaming job. If the intent was a batch load, the documentation should be updated accordingly. Given the context of "real-time event ingestion," a streaming job is likely the desired outcome.

Before:

# kafka_json_to_file_job.py
# ...
    # Read raw binary data from Kafka topic into a DataFrame
    df_raw = spark.read.format("kafka").options(**kafka_options).load()

    # ... parsing logic ...

    # Write the structured data to the local filesystem
    df_parsed.write.mode("overwrite").json(output_dir)

    spark.stop()
# ...

After (for Streaming):

# kafka_json_to_file_job.py
# ...
    # Read raw binary data from Kafka topic into a Streaming DataFrame
    # Note: 'load()' becomes 'load()' for readStream, but creates a Streaming DataFrame
    df_raw = spark.readStream.format("kafka").options(**kafka_options).load()

    # Cast raw Kafka message payload from bytes to a JSON string
    df_json = df_raw.selectExpr("CAST(value AS STRING) as json_str")

    # Parse JSON strings into structured columns based on schema
    df_parsed = df_json.select(
        from_json(col("json_str"), json_schema_str).alias("data")
    ).select("data.*")

    # Write the structured data to the local filesystem in streaming mode
    # For a learning example, a console sink or a file sink with trigger is common.
    # The 'trigger' ensures it processes data continuously/periodically.
    query = df_parsed.writeStream \
        .outputMode("append") \
        .format("json") \
        .option("path", output_dir) \
        .option("checkpointLocation", os.path.join(output_dir, "_checkpoint")) \
        .trigger(processingTime="10 seconds") \
        .start()

    # Wait for the termination of the query
    # For a real application, you might add more sophisticated shutdown hooks.
    query.awaitTermination()
    spark.stop()
# ...

Additional Notes for Streaming Implementation:

• outputMode("append"): Assumes new records are always appended. If records can be updated, update mode with a primary key might be needed.
• checkpointLocation: Crucial for streaming jobs to ensure fault tolerance and exactly-once processing. It stores progress and metadata. Should be a dedicated, persistent directory.
• trigger: Defines how often the streaming query processes new data. processingTime is suitable for micro-batch processing.
• awaitTermination(): Keeps the Spark driver alive to continuously process the stream.

2. Python Code Quality: Error Handling in json_producer.py

Issue: The json_producer.py script directly opens and reads the user_events.json file without any error handling for common file operations (e.g., FileNotFoundError).

Impact: If the input file is missing or inaccessible, the script will crash ungracefully.

Recommendation: Add a try-except block for file operations.

Before:

# json_producer.py
# ...
# Read events from input file
with open(INPUT_FILE_PATH, "r", encoding="utf-8") as f:
    records = [json.loads(line) for line in f]
# ...

After:

# json_producer.py
# ...
# Read events from input file
records = []
try:
    with open(INPUT_FILE_PATH, "r", encoding="utf-8") as f:
        for line in f:
            try:
                records.append(json.loads(line))
            except json.JSONDecodeError as e:
                print(f"Skipping malformed JSON line: {line.strip()} - Error: {e}")
except FileNotFoundError:
    print(f"Error: Input file not found at {INPUT_FILE_PATH}")
    exit(1) # Exit if the file is essential
except IOError as e:
    print(f"Error reading input file {INPUT_FILE_PATH}: {e}")
    exit(1)
# ...

3. PySpark Best Practice (Minor): Output Partitioning

Issue: The current batch write df_parsed.write.mode("overwrite").json(output_dir) will write the output JSON data into multiple part files, with the number of files determined by the number of partitions in df_parsed. For small datasets, this can lead to many small files (the "small file problem"), which can be inefficient for downstream processing.

Recommendation: For controlled file sizes, especially when writing to object storage or when the total output size is small, consider repartition() or coalesce() before writing. For this learning example, it's less critical, but good to be aware of for production systems.

Before:

    df_parsed.write.mode("overwrite").json(output_dir)

After (Example for controlled partitioning):

    # For small datasets or to control file count, consider repartitioning
    # This example repartitions to a single file, adjust as needed.
    # For a streaming job, consider .option("maxRecordsPerFile", N) if available for the sink
    df_parsed.coalesce(1).write.mode("overwrite").json(output_dir)

4. Configuration Validation (Minor)

Issue: The load_config and run_ingestion_job functions assume the input_config.yml structure is always valid (e.g., data_sources is a list, kafka_config exists within the first source, etc.). If the YAML file is malformed or missing expected keys, it could lead to KeyError or TypeError at runtime.

Recommendation: For more robust applications, add explicit checks or use dict.get() with default values to safely access configuration parameters. For a learning example, the current approach is acceptable to keep it concise.

✅ Verdict

Request Changes

The core functionality to demonstrate "Spark Streaming" is currently implemented as a batch job. Addressing this will significantly improve the alignment with the project's stated goals and provide a more accurate learning experience for real-time data ingestion. The other recommendations are primarily for robustness and best practices that can be considered for future iterations or production environments.

I've reviewed the provided PR Diff, focusing on PySpark & Data Engineering best practices, Python code quality, and potential bugs/edge cases.

🤖 AI PR Review Summary

This PR introduces a new "Ingestion" project that demonstrates real-time event processing using Kafka and Spark Structured Streaming, alongside a "Essentials" project for basic PySpark examples. It also includes necessary infrastructure changes like new dependencies, .gitignore updates, and comprehensive documentation.

The core kafka_json_to_file_job.py script is well-designed, incorporating robust configuration validation, correct Kafka integration for both streaming and batch modes, schema-based JSON parsing, and proper handling of Spark resources and checkpointing for streaming. The json_producer.py script is also well-implemented with good file and JSON parsing error handling. The accompanying documentation and Docker setup are clear and complete.

💡 Key Feedback & Recommendations

Overall, the code quality and data engineering practices are excellent. I have a few minor recommendations for small improvements, primarily around consistency, readability, and slightly enhanced robustness.

1. PySpark & Data Engineering: Batch Write Partitioning

The batch mode uses coalesce(1) which is good for small datasets to prevent too many small files. However, for a more general solution that scales with larger data, writing without coalesce(1) (allowing Spark to manage partitions) or dynamically repartitioning based on data volume might be preferred. For this specific learning project with small input, coalesce(1) is perfectly acceptable and even desirable.

Recommendation (Minor): Add a comment explaining the rationale for coalesce(1) in the context of this example, or consider removing it if the intent is to showcase default scalable behavior (though coalesce(1) is often a pragmatic choice for small results).

--- a/projects/ingestion/kafka_json_to_file_job.py
+++ b/projects/ingestion/kafka_json_to_file_job.py
@@ -102,7 +102,9 @@
                 from_json(col("json_str"), json_schema_str).alias("data")
             ).select("data.*")
 
-            # Coalesce to control partition count for small datasets
+            # Coalesce to 1 partition for small datasets to avoid many small files.
+            # For larger datasets, remove .coalesce(1) to let Spark manage partitions
+            # or repartition based on business keys.
             df_parsed.coalesce(1).write.mode("overwrite").json(output_dir)
     finally:
         spark.stop()

2. Python Code Quality: Consistent print for errors

In json_producer.py, error messages are printed to sys.stderr, which is good practice. In kafka_json_to_file_job.py, the KeyboardInterrupt message uses a regular print. While minor, using sys.stderr for all non-standard output (like errors or warnings) can help differentiate it from normal program output, especially when redirecting output.

Recommendation: Use sys.stderr for the interruption message in kafka_json_to_file_job.py.

--- a/projects/ingestion/kafka_json_to_file_job.py
+++ b/projects/ingestion/kafka_json_to_file_job.py
@@ -91,7 +91,9 @@
             try:
                 query.awaitTermination()
             except KeyboardInterrupt:
-                print("Streaming query interrupted by user. Stopping...")
+                import sys
+
+                print("Streaming query interrupted by user. Stopping...", file=sys.stderr)
         else:
             # Batch Read
             df_raw = spark.read.format("kafka").options(**kafka_options).load()

3. Python Code Quality: Type Hint for load_config

The load_config function has a docstring describing what it returns, but adding a type hint for the return value (dict) would improve readability and enable static analysis.

Recommendation: Add a return type hint to load_config.

--- a/projects/ingestion/kafka_json_to_file_job.py
+++ b/projects/ingestion/kafka_json_to_file_job.py
@@ -10,7 +10,7 @@
 
 
 # %%
-def load_config(path):
+def load_config(path: str) -> dict:
     """Loads YAML configuration file."""
     with open(path, "r", encoding="utf-8") as f:
         return yaml.safe_load(f)

4. General: requirements.txt consistency

While confluent-kafka[jsonschema] is specific to the producer and its schema handling, the pyyaml dependency is primarily used by the Spark job. It's good that these are in 3rdparty/requirements.txt, making them available globally to the Pants monorepo. This is a preference more than a strict rule, but sometimes project-specific dependencies are defined within projects/*/requirements.txt and then aggregated. For a learning monorepo, a single top-level 3rdparty/requirements.txt is simpler and fine.

Recommendation (No Change Needed, just an observation): The current approach is consistent with the existing requirements.txt structure and is acceptable for this monorepo.

✅ Verdict

Approve

The PR is excellent. The new projects are well-structured, follow best practices for both Python and PySpark, and are accompanied by comprehensive documentation and a robust Docker setup. The minor recommendations above are stylistic or for marginal robustness improvements and do not block approval. Great work!

This Pull Request introduces a new "Ingestion Project" focusing on real-time event ingestion using Kafka and Spark Streaming, alongside organizing existing examples into an "Essentials Project." The changes include updates to documentation, new dependencies, Docker Compose setup for Kafka, a Python Kafka producer, and a PySpark ingestion job.

Overall, the changes are well-structured, follow good practices, and clearly demonstrate the intended data engineering patterns.

🤖 AI PR Review Summary

This PR significantly expands the project's scope by adding a new "Ingestion Project" which showcases real-time data processing with Kafka and PySpark Structured Streaming. The new project provides a complete, runnable example, including Dockerized Kafka infrastructure, a data producer, and a robust Spark job. Documentation has been updated to reflect these additions, and necessary Python dependencies are included.

The PySpark ingestion job demonstrates excellent practices, such as externalizing configuration in YAML, robust input validation, graceful shutdown for streaming queries, and proper handling of Kafka integration with from_json parsing. The project structure is logical and easy to follow.

💡 Key Feedback & Recommendations

1. PySpark & Data Engineering Best Practices

• Robust Configuration and Error Handling: The kafka_json_to_file_job.py effectively uses a YAML configuration file and includes comprehensive validation for missing keys, which is excellent for maintainability and preventing runtime errors. The graceful shutdown for streaming jobs (KeyboardInterrupt handler) is also a strong best practice.

• Structured Streaming Checkpointing: The use of checkpointLocation for streaming writes (query.awaitTermination()) is crucial for fault tolerance and exactly-once processing, correctly implemented here.

• Batch Write coalesce(1): The comment explaining the coalesce(1) in batch mode is insightful, acknowledging its trade-offs. While acceptable for small, local examples to avoid many tiny files, it's important to remember that coalesce(1) can become a significant bottleneck in larger-scale production environments as it forces all data onto a single executor.

  Recommendation: For production-grade applications, consider removing coalesce(1) unless a specific business requirement dictates a single output file, or use repartition() with an appropriate number of partitions/keys if finer control is needed over the output file count and distribution. The current explanation in the code is sufficient for this project's scope.

2. Python Code Quality

• Docstrings and Type Hints: Functions like load_config and run_ingestion_job are well-documented with clear docstrings and type hints, enhancing readability and maintainability.

• Error Handling in Producer: The json_producer.py includes good error handling for file I/O and malformed JSON lines, making it robust against common data issues.

• Line Length: Several lines, particularly within the Spark DataFrame transformations, exceed standard line length limits (e.g., 88 characters for Black or 120 for others). While Spark chaining often leads to long lines, breaking them down can improve readability.

  Recommendation: Consider breaking long DataFrame chain operations onto multiple lines using Python's implicit line joining (parentheses) to improve readability and adhere to PEP8/Ruff guidelines.

  Before:

  df_parsed = df_json.select(
      from_json(col("json_str"), json_schema_str).alias("data")
  ).select("data.*")

  After:

  parsed_data_df = df_json.select(
      from_json(col("json_str"), json_schema_str).alias("data")
  )
  df_parsed = parsed_data_df.select("data.*")

• Producer Configuration Externalization: The json_producer.py hardcodes the Kafka bootstrap.servers and topic. While fine for a simple example, for consistency with the ingestion job's use of input_config.yml, the producer could also load its Kafka configuration from a similar YAML file or accept command-line arguments.

  Recommendation: For future enhancements or more complex scenarios, consider externalizing the Kafka configuration for json_producer.py (e.g., loading it from input_config.yml or a dedicated producer config file).

3. Bugs & Edge Cases

• No critical bugs found. The code handles common edge cases like file not found, malformed JSON lines in the producer, and graceful shutdown for streaming.
• from_json Behavior: from_json by default handles malformed JSON records by returning null for the parsed data column, rather than failing the job. This is generally desired in streaming ingestion scenarios, but for production, a dedicated Dead Letter Queue (DLQ) mechanism might be implemented to capture and reprocess these invalid records. For this example, the current behavior is acceptable.

✅ Verdict

Approve

The PR introduces valuable content and demonstrates strong data engineering principles. The code is well-written, robust, and correctly implements the described functionality. The minor recommendations focus on stylistic improvements and considerations for scaling to production, which do not hinder the immediate functionality or learning objectives of this project.